/*
**  Copyright 2010 Jim Studt. All rights reserved.
**  
**  Redistribution and use in source and binary forms, with or without modification, are
**  permitted provided that the following conditions are met:
**  
**     1. Redistributions of source code must retain the above copyright notice, this list of
**        conditions and the following disclaimer.
**  
**     2. Redistributions in binary form must reproduce the above copyright notice, this list
**        of conditions and the following disclaimer in the documentation and/or other materials
**        provided with the distribution.
**  
**  THIS SOFTWARE IS PROVIDED BY JIM STUDT ``AS IS'' AND ANY EXPRESS OR IMPLIED
**  WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
**  FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL JIM STUDT OR
**  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
**  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
**  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
**  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
**  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
**  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
**  
**  The views and conclusions contained in the software and documentation are those of the
**  authors and should not be interpreted as representing official policies, either expressed
**  or implied, of Jim Studt.
*/

#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <getopt.h>
#include <limits.h>

static int beVerbose = 0;
static const char *outputName = 0;

struct adds {
    struct adds * next;
    int appSegment;
    int length;     /* needs 2 more for the jpeg encoding */
    unsigned char *content;
} *adds = 0;

struct extracts { 
    int nth;
    char *template;
} extracts[16];

static const char short_options [] = "h?vi:x:r:o:";

static const struct option
long_options [] = {
	{ "insert",     required_argument,      NULL,           'i' },
	{ "extract",    required_argument,      NULL,           'x' },
	{ "remove",     required_argument,      NULL,           'r' },
	{ "output",     required_argument,      NULL,           'o' },
        { "help",       no_argument,            NULL,           'h' },
	{ "verbose",    no_argument,            NULL,           'v' },
        { 0, 0, 0, 0 }
};

/*
** Set and index to 1 to skip a given tag.
*/
static unsigned char skip[256] = {0};

//
// A printf to stderr, but only if in verbose mode
//
void verbose( const char *format, ...)
{
    va_list args;

    if ( !beVerbose) return;
    va_start( args, format);
    vfprintf(stderr, format, args);
    va_end(args);
}

//
// A printf to stderr and an exit(1).
//   (adds a trailing newline as well.)
//
void fatal( const char *format, ...) __attribute__ ((noreturn));
void fatal( const char *format, ...)
{
    va_list args;

    va_start( args, format);
    vfprintf(stderr, format, args);
    va_end(args);
    fputc('\n',stderr);

    exit(1);
}


static void usage(FILE *fp, int argc, char **argv)
{
    fprintf (fp,
	     "Usage: %s [options] [input file]\n"
	     " e.g.: %s -i 7=alpha0:image/png -i 7@./mask.png image.jpg\n\n"
	     "Options:\n"
	     "  -i | --insert seg@file        insert the contents of file as APPx where x=seg\n"
	     "       --insert seg=string      insert the argument as APPx where x=seg\n"
	     "  -r | --remove seg             remove all APPx segments where x=seg\n"
	     "  -x | --extract seg=template   extract all APPx segments to files named\n"
	     "                                like template with @ replaced with 0..n\n"
             "                                to disambiguate multiple occurences.\n"
	     "  -o | --output file            send output to file instead of stdout\n"
	     "  -v | --verbose                Print a lot of debug messages\n"
	     "  -h | --help                   This help message\n"
	     "",
	     argv[0], argv[0]);
}

static void do_options(int argc, char **argv)
{
    for (;;) {
	int index;
	int c;
                
	c = getopt_long (argc, argv, short_options, long_options, &index);

	if (-1 == c) break;

	switch (c) {
	  case 0: /* getopt_long() flag */
	    break;
	  case 'o':
	    outputName = strdup(optarg);
	    verbose("output will go to %s\n", outputName);
	    break;
	  case 'r':
	      {
		  char *p;
		  int v = strtol( optarg, &p, 10);
		  if ( *p != 0 || v < 0 || v > 15) fatal("Invalid argument to --remove, '%s'", optarg);
		  verbose("will skip APP%d\n", v);
		  skip[(unsigned char)(0xe0 + v)] = 1;
	      }
	      break;
	  case 'x':
	      {
		  int seg;
		  char templ[65536];

		  if ( sscanf( optarg, "%d=%65535s", &seg, templ) < 2) {
		      fatal("Poorly formed --extract: %s\n", optarg);
		  }

		  if ( seg < 0 || seg > 15) fatal("Segment number of an extract needs to be from 0 to 15.");

		  extracts[seg].template = strdup(templ);
	      }
	      break;
	  case 'i':
	      {
		  int seg;
		  char arg[65536] = "";
		  int contentLength = 0;
		  void *contentPointer = 0;
		  struct adds **tail = &adds;

		  if ( sscanf( optarg, "%d=%65535s", &seg, arg) < 2) {
		      if ( sscanf( optarg, "%d@%65535s", &seg, arg) < 2) {
			  fatal("Poorly formed insertion.");
		      }
		      /* we have a @filename style, get the content */
		      contentPointer = malloc(65536); /* maximum encodable */
		      {
			  int len;
			  FILE *f = fopen( arg, "rb");
			  if ( !f) fatal("Unable to open %s for insertion: %s", arg, strerror(errno));
			  len = fread( contentPointer, 1, 65534, f);
			  if ( ferror(f)) fatal("Error reading %s for insertion: %s", arg, strerror(errno));
			  if ( !feof(f)) fatal("Truncated %s for insertion, 65534 bytes maximum.", arg);
			  fclose(f);
			  contentPointer = realloc( contentPointer, len);
			  contentLength = len;
		      }
		  } else {
		      /* we have a literal style */
		      contentPointer = strdup(arg);
		      contentLength = strlen(arg);
		  }
		  if ( seg < 0 || seg > 15) fatal("Segment number of an insert needs to be from 0 to 15.");

		  while( *tail != 0) tail = & (*tail)->next;

		  *tail = (struct adds *)malloc(sizeof(struct adds));
		  (*tail)->next = 0;
		  (*tail)->appSegment = seg;
		  (*tail)->length = contentLength;
		  (*tail)->content = contentPointer;
	      }
	      break;
	  case 'h':
	  case '?':
	    usage (stdout, argc, argv);
	    exit (EXIT_SUCCESS);
	  case 'v':
	    beVerbose++;
	    break;
	  default:
	    usage (stderr, argc, argv);
	    exit (EXIT_FAILURE);
	}
    }
}

static void putOut( int o, FILE *f) {
    if ( fputc( o, f) == EOF) fatal("Failed writing output: %s", strerror(errno));
}

static int doEntropy( FILE *in, FILE *out, int c1, int c2)
{
    int count = 2;

    putOut( c1, out);
    putOut( c2, out);
    /* copy entropy encoded data */
    for (;;) {
	int ch = fgetc(in);
	if ( ch == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
	
	if ( ch == 0xff) {
	    ch = fgetc(in);
	    if ( ch == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
	    if ( ch == 0) {
		putOut(0xff,out);
		putOut(0,out);
		count += 2;
	    } else {
		verbose("entropy encoded section, %d bytes\n", count);
		return ch;
	    }
	} else {
	    putOut(ch,out);
	    count++;
	}
    }
}

/*
** We learn the structure of a JPEG from http://www.w3.org/Graphics/JPEG/itu-t81.pdf
** Specifically Appendix B.
** In short we have markers of the from 0xff 0x?? followed by a two byte count, except
** where they are not followed by a two byte count, but instead by an stream of bytes with 
** their 0xff bytes escaped to 0xff 0x00. Oh, and d0-d7, d8, d9, and 01 which don't have any
** content.
*/
void processJPEG( FILE *in, FILE *out)
{
    int shouldEnd = 0;
    int needInsert = (adds != 0);

    for (;;) {
	int tag;
	int ff = fgetc(stdin);
	if ( ff == EOF) {
	    if ( ferror(stdin)) fatal("Failed reading input: %s", strerror(errno));
	    if ( !shouldEnd) fatal("Premature end of JPEG input");
	    break;  /* This is our loop exit <<<<<<<<<<<<< */
	}
	
	tag = fgetc(stdin);
	if ( tag == EOF) fatal("Premature end of JPEG input");
	
	if ( ff != 0xff || tag == 0) tag = doEntropy( in,out, ff, tag);

	if ( (tag >= 0xd0 && tag <= 0xd9) || tag == 0x00) {
	    /* a tag with no content, marker only */
	    verbose("copied xff%02x\n", tag);
	    putOut( 0xff, out);
	    putOut( tag, out);
	    if ( tag == 0xd9) shouldEnd = 1;
	} else {
	    /* copy the marker's block */
	    int msb = fgetc(stdin);
	    int lsb = fgetc(stdin);
	    int len;
	    int skipping = skip[(unsigned char)tag];
	    FILE *extractFile = 0;

	    if ( lsb == EOF || msb == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
	    len = msb*256+lsb;

	    if ( needInsert && tag == 0xda) {
		/* do our insert right before the SOS tag (start of scan), no particular reason, but legal */
		struct adds *a = adds;

		for ( a = adds; a; a = a->next) {
		    int i;
		    int len = a->length;
		    putOut( 0xff, out);
		    putOut( 0xe0 + a->appSegment, out);
		    len += 2;
		    putOut( (len >> 8) & 0xff, out);
		    putOut( len & 0xff, out);
		    len -= 2;
		    for ( i = 0; i < len; i++) putOut( a->content[i], out);
		    verbose("inserted a APP%d length %d bytes\n", a->appSegment, len);
		}
		needInsert = 0;
	    }

	    if ( tag >= 0xe0 && tag <= 0xef ) {
		int seg = ( tag & 0x0f );
		if ( seg && extracts[seg].template) {
		    char name[PATH_MAX];
		    char *n = name;
		    char *lim = name+PATH_MAX-1;
		    char *p;
		    
		    for ( p = extracts[seg].template; *p; p++) {
			if ( *p != '@') {
			    if ( n < lim) *n++ = *p;
			} else {
			    int cnt = snprintf( n, lim-n-1, "%d", extracts[seg].nth++);
			    n += cnt;
			}
		    }
		    *n = 0;

		    verbose("extracting APP%d(%d bytes) to %s\n", seg, len, name);
		    extractFile = fopen(name,"wb");
		    if ( !extractFile) fatal("Failed to open %s for APP%d data: %s", name, seg, strerror(errno));
		}
	    }

	    if ( !skipping) {
		verbose("copied xff%02x %d bytes\n", tag, len);
		putOut( 0xff, out);
		putOut( tag, out);
		putOut( msb, out);
		putOut( lsb, out);
	    } else {
		verbose("skipped xff%02x %d bytes\n", tag, len);
	    }
	    
	    /* we take 2 off for the len field itself */
	    for ( len -= 2; len > 0; len--) {
		int ch = fgetc(stdin);
		if ( ch == EOF) fatal("Permature end of JPEG input: %s", strerror(errno));
		if ( extractFile) {
		    if ( fputc( ch, extractFile) == EOF) {
			fatal("Failed to write to extracted data segment: %s", strerror(errno));
		    }
		}
		if ( !skipping) putOut(ch,out);
	    }
	    if ( extractFile) {
		if ( fclose(extractFile) == EOF) fatal("Failed to close extract file: %s", strerror(errno));
		extractFile = 0;
	    }
	}		
    }
}

int main( int argc, char **argv)
{
    do_options(argc, argv);

    /*
    ** Reopen stdin as required
    */
    switch( argc - optind) {
      case 0:
	/* using stdin */
	break;
      case 1:
	  {
	      const char *inFile = argv[optind];
	      stdin = freopen( inFile, "r", stdin);
	      if ( stdin == 0) fatal("Unable to open %s for reading: %s", inFile, strerror(errno));
	      verbose("input is reopened to %s\n", inFile);
	  }
	break;
      default:
	fatal("Too many arguments, can only process one file.");
    }

    /*
    ** Reopen stdout as required
    */ 
    {
	FILE *outFile = stdout;

	if ( outputName) {
	    /*
	    ** I wonder if maybe I should write to a temporary file and rename after success.
	    ** I don't expect to fail, but that would allow input and output to be the same file.
	    ** Maybe later.
	    */
	    outFile = fopen( outputName, "wb");
	    if ( outFile == 0) fatal("Unable to write to %s: %s", outputName, strerror(errno));
	    verbose("output is reopened to %s\n", outputName);
	}
	
	processJPEG( stdin, outFile);
	
	if ( fclose(outFile) == EOF) fatal("Failed to close output file: %s", strerror(errno));
    }

    if ( fclose(stdin) == EOF) fatal("Failed to close input file: %s", strerror(errno));
    
    return 0;
}
